{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['HTTP_PROXY'] = 'http://127.0.0.1:7890'\n",
    "os.environ['HTTPS_PROXY'] = 'http://127.0.0.1:7890'\n",
    "\n",
    "from langchain import PromptTemplate, FewShotPromptTemplate, LLMChain, OpenAI\n",
    "from langchain_core.messages import HumanMessage, SystemMessage\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "parser = StrOutputParser()\n",
    "\n",
    "model = ChatOpenAI(model=\"gpt-4o-mini\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import ChatPromptTemplate, FewShotChatMessagePromptTemplate\n",
    "\n",
    "examples = [\n",
    "    {\"input\": \"2 🦜 2\", \"output\": \"4\"},\n",
    "    {\"input\": \"2 🦜 3\", \"output\": \"5\"},\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[HumanMessage(content='2 🦜 2'), AIMessage(content='4'), HumanMessage(content='2 🦜 3'), AIMessage(content='5')]\n"
     ]
    }
   ],
   "source": [
    "# This is a prompt template used to format each individual example.\n",
    "example_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"human\", \"{input}\"),\n",
    "        (\"ai\", \"{output}\"),\n",
    "    ]\n",
    ")\n",
    "few_shot_prompt = FewShotChatMessagePromptTemplate(\n",
    "    example_prompt=example_prompt,\n",
    "    examples=examples,\n",
    "    # suffix=\"Human: {input}\",\n",
    ")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[HumanMessage(content='2 🦜 2'),\n",
       " AIMessage(content='4'),\n",
       " HumanMessage(content='2 🦜 3'),\n",
       " AIMessage(content='5')]"
      ]
     },
     "execution_count": 93,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "few_shot_prompt.invoke({}).messages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Human: 2 🦜 2\\nAI: 4\\nHuman: 2 🦜 3\\nAI: 5'"
      ]
     },
     "execution_count": 86,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "few_shot_prompt.format()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(few_shot_prompt.invoke({}).to_messages())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content='Based on the previous pattern, the 🦜 operation appears to be addition. Therefore:\\n\\n\\\\[ 3 🦜 3 = 3 + 3 = 6 \\\\]', response_metadata={'token_usage': {'completion_tokens': 37, 'prompt_tokens': 30, 'total_tokens': 67}, 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': '', 'finish_reason': 'stop', 'logprobs': None}, id='run-1e976e57-ef4f-4407-9423-2ee339e6352c-0', usage_metadata={'input_tokens': 30, 'output_tokens': 37, 'total_tokens': 67})"
      ]
     },
     "execution_count": 94,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "final_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", \"You are a wondrous wizard of math.\"),\n",
    "        few_shot_prompt,\n",
    "        (\"human\", \"{input}\"),\n",
    "    ]\n",
    ")\n",
    "# chain = model | final_prompt\n",
    "chain = final_prompt | model\n",
    "\n",
    "chain.invoke({\"input\": \"What's 3 🦜 3?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"System: You are a wondrous wizard of math.\\nHuman: 2 🦜 2\\nAI: 4\\nHuman: 2 🦜 3\\nAI: 5\\nHuman: What's 3 🦜 3?\""
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "final_prompt.format(input=\"What's 3 🦜 3?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "System: You are a wondrous wizard of math.\n",
      "Human: 2 🦜 2\n",
      "AI: 4\n",
      "Human: 2 🦜 3\n",
      "AI: 5\n",
      "Human: What's 3 🦜 3?\n"
     ]
    }
   ],
   "source": [
    "print(final_prompt.format(input=\"What's 3 🦜 3?\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import PromptTemplate\n",
    "\n",
    "example_prompt = PromptTemplate.from_template(\"Question: {question}\\n{answer}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "qa_examples = [\n",
    "    {\n",
    "        \"question\": \"Who lived longer, Muhammad Ali or Alan Turing?\",\n",
    "        \"answer\": \"\"\"\n",
    "            Are follow up questions needed here: Yes.\n",
    "            Follow up: How old was Muhammad Ali when he died?\n",
    "            Intermediate answer: Muhammad Ali was 74 years old when he died.\n",
    "            Follow up: How old was Alan Turing when he died?\n",
    "            Intermediate answer: Alan Turing was 41 years old when he died.\n",
    "            So the final answer is: Muhammad Ali\n",
    "            \"\"\",\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"When was the founder of craigslist born?\",\n",
    "        \"answer\": \"\"\"\n",
    "            Are follow up questions needed here: Yes.\n",
    "            Follow up: Who was the founder of craigslist?\n",
    "            Intermediate answer: Craigslist was founded by Craig Newmark.\n",
    "            Follow up: When was Craig Newmark born?\n",
    "            Intermediate answer: Craig Newmark was born on December 6, 1952.\n",
    "            So the final answer is: December 6, 1952\n",
    "            \"\"\",\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"Who was the maternal grandfather of George Washington?\",\n",
    "        \"answer\": \"\"\"\n",
    "            Are follow up questions needed here: Yes.\n",
    "            Follow up: Who was the mother of George Washington?\n",
    "            Intermediate answer: The mother of George Washington was Mary Ball Washington.\n",
    "            Follow up: Who was the father of Mary Ball Washington?\n",
    "            Intermediate answer: The father of Mary Ball Washington was Joseph Ball.\n",
    "            So the final answer is: Joseph Ball\n",
    "            \"\"\",\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"Are both the directors of Jaws and Casino Royale from the same country?\",\n",
    "        \"answer\": \"\"\"\n",
    "            Are follow up questions needed here: Yes.\n",
    "            Follow up: Who is the director of Jaws?\n",
    "            Intermediate Answer: The director of Jaws is Steven Spielberg.\n",
    "            Follow up: Where is Steven Spielberg from?\n",
    "            Intermediate Answer: The United States.\n",
    "            Follow up: Who is the director of Casino Royale?\n",
    "            Intermediate Answer: The director of Casino Royale is Martin Campbell.\n",
    "            Follow up: Where is Martin Campbell from?\n",
    "            Intermediate Answer: New Zealand.\n",
    "            So the final answer is: No\n",
    "            \"\"\",\n",
    "    },\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Question: Who lived longer, Muhammad Ali or Alan Turing?\n",
      "\n",
      "            Are follow up questions needed here: Yes.\n",
      "            Follow up: How old was Muhammad Ali when he died?\n",
      "            Intermediate answer: Muhammad Ali was 74 years old when he died.\n",
      "            Follow up: How old was Alan Turing when he died?\n",
      "            Intermediate answer: Alan Turing was 41 years old when he died.\n",
      "            So the final answer is: Muhammad Ali\n",
      "            \n"
     ]
    }
   ],
   "source": [
    "print(example_prompt.invoke(qa_examples[0]).text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Question: Who lived longer, Muhammad Ali or Alan Turing?\n",
      "\n",
      "            Are follow up questions needed here: Yes.\n",
      "            Follow up: How old was Muhammad Ali when he died?\n",
      "            Intermediate answer: Muhammad Ali was 74 years old when he died.\n",
      "            Follow up: How old was Alan Turing when he died?\n",
      "            Intermediate answer: Alan Turing was 41 years old when he died.\n",
      "            So the final answer is: Muhammad Ali\n",
      "            \n",
      "\n",
      "Question: When was the founder of craigslist born?\n",
      "\n",
      "            Are follow up questions needed here: Yes.\n",
      "            Follow up: Who was the founder of craigslist?\n",
      "            Intermediate answer: Craigslist was founded by Craig Newmark.\n",
      "            Follow up: When was Craig Newmark born?\n",
      "            Intermediate answer: Craig Newmark was born on December 6, 1952.\n",
      "            So the final answer is: December 6, 1952\n",
      "            \n",
      "\n",
      "Question: Who was the maternal grandfather of George Washington?\n",
      "\n",
      "            Are follow up questions needed here: Yes.\n",
      "            Follow up: Who was the mother of George Washington?\n",
      "            Intermediate answer: The mother of George Washington was Mary Ball Washington.\n",
      "            Follow up: Who was the father of Mary Ball Washington?\n",
      "            Intermediate answer: The father of Mary Ball Washington was Joseph Ball.\n",
      "            So the final answer is: Joseph Ball\n",
      "            \n",
      "\n",
      "Question: Are both the directors of Jaws and Casino Royale from the same country?\n",
      "\n",
      "            Are follow up questions needed here: Yes.\n",
      "            Follow up: Who is the director of Jaws?\n",
      "            Intermediate Answer: The director of Jaws is Steven Spielberg.\n",
      "            Follow up: Where is Steven Spielberg from?\n",
      "            Intermediate Answer: The United States.\n",
      "            Follow up: Who is the director of Casino Royale?\n",
      "            Intermediate Answer: The director of Casino Royale is Martin Campbell.\n",
      "            Follow up: Where is Martin Campbell from?\n",
      "            Intermediate Answer: New Zealand.\n",
      "            So the final answer is: No\n",
      "            \n",
      "\n",
      "Question: Who was the father of Mary Ball Washington?\n"
     ]
    }
   ],
   "source": [
    "from langchain_core.prompts import FewShotPromptTemplate\n",
    "\n",
    "prompt = FewShotPromptTemplate(\n",
    "    examples=qa_examples,\n",
    "    example_prompt=example_prompt,\n",
    "    # prefix=\"You are a helpful assistant.\",\n",
    "    suffix=\"Question: {input}\",\n",
    "    input_variables=[\"input\"],\n",
    ")\n",
    "\n",
    "print(\n",
    "    prompt.invoke({\"input\": \"Who was the father of Mary Ball Washington?\"}).to_string()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_chroma import Chroma\n",
    "from langchain_core.example_selectors import SemanticSimilarityExampleSelector\n",
    "from langchain_openai import OpenAIEmbeddings\n",
    "\n",
    "example_selector = SemanticSimilarityExampleSelector.from_examples(\n",
    "    # This is the list of examples available to select from.\n",
    "    qa_examples,\n",
    "    # This is the embedding class used to produce embeddings which are used to measure semantic similarity.\n",
    "    OpenAIEmbeddings(),\n",
    "    # This is the VectorStore class that is used to store the embeddings and do a similarity search over.\n",
    "    Chroma,\n",
    "    # This is the number of examples to produce.\n",
    "    k=1,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Examples most similar to the input: Who was the father of Mary Ball Washington?\n",
      "\n",
      "\n",
      "【\n",
      "answer: \n",
      "            Are follow up questions needed here: Yes.\n",
      "            Follow up: Who was the mother of George Washington?\n",
      "            Intermediate answer: The mother of George Washington was Mary Ball Washington.\n",
      "            Follow up: Who was the father of Mary Ball Washington?\n",
      "            Intermediate answer: The father of Mary Ball Washington was Joseph Ball.\n",
      "            So the final answer is: Joseph Ball\n",
      "            \n",
      "question: Who was the maternal grandfather of George Washington?\n",
      "】\n"
     ]
    }
   ],
   "source": [
    "# Select the most similar example to the input.\n",
    "question = \"Who was the father of Mary Ball Washington?\"\n",
    "selected_examples = example_selector.select_examples({\"question\": question})\n",
    "print(f\"Examples most similar to the input: {question}\")\n",
    "for example in selected_examples:\n",
    "    print(\"\\n\")\n",
    "    print('【')\n",
    "    for k, v in example.items():\n",
    "        print(f\"{k}: {v}\")\n",
    "    print('】')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'answer': '\\n            Are follow up questions needed here: Yes.\\n            Follow up: Who was the mother of George Washington?\\n            Intermediate answer: The mother of George Washington was Mary Ball Washington.\\n            Follow up: Who was the father of Mary Ball Washington?\\n            Intermediate answer: The father of Mary Ball Washington was Joseph Ball.\\n            So the final answer is: Joseph Ball\\n            ',\n",
       "  'question': 'Who was the maternal grandfather of George Washington?'}]"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "selected_examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Question: Who was the maternal grandfather of George Washington?\n",
      "\n",
      "            Are follow up questions needed here: Yes.\n",
      "            Follow up: Who was the mother of George Washington?\n",
      "            Intermediate answer: The mother of George Washington was Mary Ball Washington.\n",
      "            Follow up: Who was the father of Mary Ball Washington?\n",
      "            Intermediate answer: The father of Mary Ball Washington was Joseph Ball.\n",
      "            So the final answer is: Joseph Ball\n",
      "            \n",
      "\n",
      "Question: Who was the father of Mary Ball Washington?\n"
     ]
    }
   ],
   "source": [
    "prompt = FewShotPromptTemplate(\n",
    "    example_selector=example_selector,\n",
    "    example_prompt=example_prompt,\n",
    "    suffix=\"Question: {input}\",\n",
    "    input_variables=[\"input\"],\n",
    ")\n",
    "\n",
    "print(\n",
    "    prompt.invoke({\"input\": \"Who was the father of Mary Ball Washington?\"}).to_string()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content='The father of Mary Ball Washington was Joseph Ball.', response_metadata={'token_usage': {'completion_tokens': 10, 'prompt_tokens': 103, 'total_tokens': 113}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_0f03d4f0ee', 'finish_reason': 'stop', 'logprobs': None}, id='run-ae96f9c7-ac89-47ba-8074-69197b89bef5-0', usage_metadata={'input_tokens': 103, 'output_tokens': 10, 'total_tokens': 113})"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain = prompt | model\n",
    "chain.invoke({\"input\": \"Who was the father of Mary Ball Washington?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "agent",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
